import torch
# from src.tensor_transform import *
device = torch.device('cuda')

def get_error(H, H1, C_hard, C_soft, voxel, output, Ke_hard, Fe_hard, Ke_soft, Fe_soft, jacobi):
    """computer batch error between predicated homogenized tensor and the groundtruth

    Args:
        H (object): Homo 3D full instance\
        H1 (object): Homo 3D instance
        C (torch.cuda.FloatTensor): base material elastic tensor, it should be batch_size*36
        voxel (torch.cuda.FloatTensor):  it should be batch_size*36*N*N*N
        output (torch.cuda.FloatTensor): predicated displacement,it should be batch_size*18*N*N*N
        transform (torch.cuda.FloatTensor): Jacobian transformer, it should be batch_size*3*3
        Ke (torch.cuda.FloatTensor): element stiffness matrics, it should be batch_size*24*24
        Fe (torch.cuda.FloatTensor): element force matrics, it should be batch_size*24*6

    Returns:
        float: all_error
        torch.cuda.FloatTensor: batch_error, it should be batch_size
    """

    batch_predict=torch.zeros_like(C_hard.contiguous().view(-1,6,6))
    batch_actual=torch.zeros_like(C_hard.contiguous().view(-1,6,6))

    batch_error=torch.zeros(C_hard.shape[0],device=C_hard.device)

    for i in range(C_hard.shape[0]):

        C_homo=H.homo_new(voxel[i],output[i],Ke_hard[i],Fe_hard[i])
        # C_final=tensor_transform(C_homo,jacobi[i])
        batch_predict[i]=C_homo

        U=H1.solve_by_torch(voxel[i],Ke_hard[i],Fe_hard[i], maxit = 5000)
        C_homo=H.homo_new(voxel[i],U,Ke_hard[i],Fe_hard[i])
        # C_final=tensor_transform(C_homo,jacobi[i])
        batch_actual[i]=C_homo
        #torch.cuda.empty_cache()
        
        batch_error[i]=torch.linalg.norm(batch_actual[i]-batch_predict[i])/torch.linalg.norm(batch_actual[i])
    
    return batch_error.mean(),batch_error